import psycopg2
import pandas as pd
pg_host = "127.0.0.1"
pg_port = '5432'
pg_user = "postgres"
pg_password = "Pass2word"
pg_database = "blog"

# conn_string = "host=" + pg_host + " port=" + pg_port + " dbname=" + pg_database + " user=" + pg_user + " password=" + pg_password
# conn = psycopg2.connect(conn_string)
# curs = conn.cursor()
#
# sql_command = "select tablename from pg_tables where schemaname='public' "
# sql_command1 = "select * from authors"
# curs.execute(sql_command1)
# print(curs.fetchall())

# data = pd.DataFrame(curs.fetchall(), columns=select_columns)


class PgConn:
    def __init__(self, pg_host, pg_user, pg_password, pg_database, pg_port='5432'):
        self.conn = psycopg2.connect(database=pg_database, user=pg_user, password=pg_password, host=pg_host, port=pg_port)
        self.cursor = self.conn.cursor()

    def get_tables(self, schemaname='public'):
        """
        查询所有表明
        :param schemaname: 默认public
        :return:
        """
        # self.cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = '%s'" % schemaname)
        self.cursor.execute(f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{schemaname}'")
        return self.cursor.fetchall()

    def get_database(self):
        """
        查看所有数据库
        :return:
        """
        self.cursor.execute("select datname from pg_database;")
        return self.cursor.fetchall()

    def get_column_info(self, table_name, schemaname='public'):
        """
        查询表的列名
        :param table_name:
        :return: list of (schema, table colume_name, index, xx, xx, data_type, max_length)
        """
        self.cursor.execute(f"SELECT * FROM information_schema.columns WHERE table_schema = '{schemaname}' and table_name ='{table_name}'")
        return self.cursor.fetchall()

    def get_table_data(self, table_name, schemaname='public'):
        """
        查询表信息
        :param table_name:
        :return:
        """
        sql_command = f"SELECT * FROM {schemaname}.{table_name}"
        df_chunk = pd.read_sql(sql_command, self.conn, chunksize=1000)
        df_data = pd.concat(df_chunk)
        df_data.to_csv(f'pg_{table_name}.csv', sep=',', header=True, index=False)
        return df_data


if __name__ == '__main__':
    pg1 = PgConn(pg_host, pg_user, pg_password, pg_database, pg_port)
    print(pg1.get_tables())
    print(pg1.get_database())
    for i in pg1.get_column_info('authors'):
        # print(i)
        print(i[3], i[5], i[7], i[8])
    print(pg1.get_table_data('entries'))
